from ultralytics import YOLO


def main(args):
    model = YOLO(args.src)

    kwargs = {}

    if args.unit_size == 32:
        pass
    elif args.unit_size == 16:
        kwargs['half'] = True
    elif args.unit_size == 8:
        kwargs['int8'] = True
    else:
        raise ValueError('Invalid unit size: %d' % args.unit_size)

    if args.imgsz is not None:
        imgsz = [int(x) for x in args.imgsz.split(',')]
        if len(imgsz) == 2:
            pass
        elif len(imgsz == 1):
            imgsz = imgsz * 2
        else:
            imgsz = imgsz[0:2]
        kwargs['imgsz'] = imgsz

    if args.format == 'onnx':
        # dynamic 只支持 CPU
        # kwargs['dynamic'] = True

        # kwargs['simplify'] = True
        pass

    # 导出模型
    result = model.export(
        **kwargs,
        format=args.format)
    print(result)


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('src', type=str)
    parser.add_argument('--format', type=str, default='onnx')
    parser.add_argument('--unit-size', type=int, default=32)
    parser.add_argument('--imgsz', type=str, default='640,640')
    args = parser.parse_args()
    main(args)
